package com.tdk.mybatisplus.demo.common.util;

import com.alibaba.excel.enums.CellDataTypeEnum;
import com.alibaba.excel.metadata.Head;
import com.alibaba.excel.metadata.data.WriteCellData;
import com.alibaba.excel.util.MapUtils;
import com.alibaba.excel.write.metadata.holder.WriteSheetHolder;
import com.alibaba.excel.write.style.column.AbstractColumnWidthStyleStrategy;
import java.io.UnsupportedEncodingException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import lombok.SneakyThrows;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.poi.ss.usermodel.Cell;

/**
 * <p>
 * 设置自适应列宽
 * </p>
 *
 * @author: taodingkai
 * @modified:
 * @since: 2022/5/24 10:00
 */
public class CustomColumnWidthStyleStrategy extends AbstractColumnWidthStyleStrategy {

  private static final int                                 MAX_COLUMN_WIDTH = 30;
  private static final int                                 MIN_COLUMN_WIDTH = 20;
  public static final  String                              CHARSET_NAME_GBK = "GBK";
  private final        Map<Integer, Map<Integer, Integer>> CACHE            = MapUtils.newHashMapWithExpectedSize(8);

  @SneakyThrows
  @Override
  protected void setColumnWidth(
      WriteSheetHolder writeSheetHolder, List<WriteCellData<?>> cellDataList, Cell cell, Head head,
      Integer relativeRowIndex, Boolean isHead) {
    boolean needSetWidth = isHead || !CollectionUtils.isEmpty(cellDataList);
    if (needSetWidth) {
      Map<Integer, Integer> maxColumnWidthMap = CACHE
          .computeIfAbsent(writeSheetHolder.getSheetNo(), (key) -> new HashMap(16));
      Integer               columnWidth       = this.dataLength(cellDataList, cell, isHead);
      if (columnWidth > MAX_COLUMN_WIDTH) {
        columnWidth = MAX_COLUMN_WIDTH;
      }
      if (columnWidth < MIN_COLUMN_WIDTH) {
        columnWidth = MIN_COLUMN_WIDTH;
      }
      Integer columnWidthInMap = maxColumnWidthMap.get(cell.getColumnIndex());
      if (columnWidthInMap == null) {
        maxColumnWidthMap.put(cell.getColumnIndex(), columnWidth);
        writeSheetHolder.getSheet().setColumnWidth(cell.getColumnIndex(), columnWidth * 256);
      }
    }
  }

  private Integer dataLength(List<WriteCellData<?>> cellDataList, Cell cell, Boolean isHead)
      throws UnsupportedEncodingException {

    if (isHead) {
      return cell.getStringCellValue().getBytes(CHARSET_NAME_GBK).length;
    } else {
      WriteCellData<?> cellData = (WriteCellData) cellDataList.get(0);
      CellDataTypeEnum type     = cellData.getType();
      if (type == null) {
        return -1;
      } else {
        switch (type) {
          case STRING:
            return cellData.getStringValue().getBytes(CHARSET_NAME_GBK).length;
          case BOOLEAN:
            return cellData.getBooleanValue().toString().getBytes(CHARSET_NAME_GBK).length;
          case NUMBER:
            return cellData.getNumberValue().toString().getBytes(CHARSET_NAME_GBK).length;
          default:
            return -1;
        }
      }
    }
  }
}